import gradio as gr
from PIL import Image
import numpy as np
from sklearn.cluster import KMeans

def extract_colors(image):
    # 将PIL图像转换为numpy数组
    image = image.convert("RGB")
    image = image.resize((100, 100))  # 缩小图像以提高处理速度
    pixels = np.array(image).reshape((-1, 3))

    # 使用KMeans聚类算法找到5个主要颜色
    kmeans = KMeans(n_clusters=5)
    kmeans.fit(pixels)
    
    # 获取聚类中心的颜色值
    colors = kmeans.cluster_centers_.astype(int)
    hex_colors = ['#%02x%02x%02x' % (int(r), int(g), int(b)) for r, g, b in colors]
    
    return hex_colors

def display_colors(image):
    colors = extract_colors(image)
    # 创建一个带有颜色块和十六进制颜色值的HTML字符串
    color_blocks = ''.join([f'<div style="background-color:{color}; display:inline-block; width:50px; height:50px;"></div>' for color in colors])
    color_hexes = '<br>'.join(colors)
    return f'{color_blocks}<br>{color_hexes}'

# 创建Gradio接口
iface = gr.Interface(
    fn=display_colors,
    inputs=gr.inputs.Image(type='pil'),
    outputs=gr.outputs.HTML(),
    title="图像主色调提取",
    description="上传一张图片，分析图片的主色调调色板（5个颜色），并显示每个颜色的十六进制表示形式。"
)

# 启动Gradio应用
iface.launch()